import matplotlib.pyplot as plt
import numpy as np

class coeff_strut:
    def __init__(self,coeff,pred_t):
        self.coeff=coeff;
        self.predict_t=pred_t;
    def get_coeff(self):
        return self.coeff
    def get_predt(self):
        return self.predict_t


# target location
target_x=20.5;
target_y=20;
target_vx=0;
target_vy=0;
target_ax=0;
target_ay=0;

# ego original state
ego_x=20;
ego_y=0;
ego_vx=-0.8;# m/s
ego_vy=16;
ego_ax=0;
ego_ay=-0.5;

# trans ego state to reference line
delta_y=target_y-ego_y
delta_x=target_x-ego_x
thirdline=(delta_x**2+delta_y**2)**0.5
cos_theta=delta_y/thirdline;
sin_theta=delta_x/thirdline
# lateral velocity and accelerate
ego_vx2refline=ego_vx*cos_theta-ego_vy*sin_theta
ego_ax2refline=ego_ax*cos_theta-ego_ay*sin_theta
ego_vy2refline=ego_vy*cos_theta+ego_vx*sin_theta
ego_ay2refline=ego_ay*cos_theta+ego_ax*sin_theta

origion_state=[0,ego_vx2refline,ego_ax2refline]
origion_longitude_state=[ego_vy2refline,ego_ay2refline]

# lateral planning quintic polynomial(5)
d_array=np.arange(-0.5,-5,-0.5)
finalstate_array=[]
for i in d_array:
    finalstate_array.append([i,0,0])
steptime_array=np.arange(0.5,5,0.5)
coefficient_ar=[]
for final_state in finalstate_array:
    for t in steptime_array:
        O=origion_state+final_state
        M=np.array([
            [1,0,0,0,0,0],
            [0,1,0,0,0,0],
            [0,0,2,0,0,0],
            [1,t,t**2,t**3,t**4,t**5],
            [0,1,2*t,3*t**2,4*t**3,5*t**4],
            [0,0,2,6*t,12*t**2,20*t**3]
            ])
        O=np.array([O]).T
        t_co=np.dot(np.linalg.inv(M),O)
        coefficient_ar.append(coeff_strut(t_co,t))
# select lateral traj

# lateral planning res
plt.figure(1)
for pla_info in coefficient_ar:
    t_coe=np.array([pla_info.get_coeff()]).T.flatten()
    t_coe=np.flipud(t_coe)
    func=np.poly1d(t_coe.astype(float))
    x=np.linspace(0,pla_info.get_predt(),30)
    y=func(x)

    plt.plot(x,y)
plt.title(" lateral planning set ")
plt.xlabel("predict time [s]")
plt.ylabel("lateral distance [m]")


# longitude palnning quartic polynomial(4)
step_num=20
v_array=np.linspace(13,25,step_num)
a_array=np.linspace(0,0,step_num)
finalstate_array=[]
for i in range(step_num):
    finalstate_array.append([v_array[i],0])
coefficient_longitude_ar=[]
for final_longitude_state in finalstate_array:
    for t in steptime_array:
        O = origion_longitude_state + final_longitude_state
        M=[
            [1,0,0,0],
            [0,2,0,0],
            [1, 2*t,3*t**2,4*t**3],
            [0,2,6*t,12*t**2]
        ]
        O = np.array([O]).T
        t_co = np.dot(np.linalg.inv(M), O)
        coefficient_longitude_ar.append(coeff_strut(t_co,t))

plt.figure(2)
plt.title("longitude planning")
for pla_info in coefficient_longitude_ar:
    t_coe = np.array([pla_info.get_coeff()]).T.flatten()
    t_coe = np.flipud(t_coe)
    t_coe[0]=4*t_coe[0]
    t_coe[1]=3*t_coe[1]
    t_coe[2]=2*t_coe[2]
    func = np.poly1d(t_coe.astype(float))
    x = np.linspace(0, pla_info.get_predt(), 30)
    y = func(x)
    plt.plot(x, y)
plt.title(" longitude planning set ")
plt.xlabel("predict time [s]")
plt.ylabel("longitude speed [m/s]")

# show res
plt.figure(4)
plt.title(" ref location of 2 vehicles ")
plt.plot(target_x,target_y,"ro")
plt.plot(ego_x,ego_y,"ro")

plt.show()
